"""
混合检索器 - 融合向量检索、关键词检索和图检索
实现多模态检索策略，提高检索质量和覆盖度
"""

import asyncio
from typing import List, Dict, Any, Optional, Union
import re
from collections import Counter
import math

from ...core.interfaces import RetrievalInterface, SearchResult, Document
from ...core.events import EventEmitter
from ...core.container import container
from ...data_layer.data_manager import DataManager
from ..ingestion.embedding_generator import EmbeddingGenerator


class BM25Retriever:
    """BM25关键词检索器"""
    
    def __init__(self, k1: float = 1.2, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self.documents: List[Document] = []
        self.doc_freqs: List[Dict[str, int]] = []
        self.idf_values: Dict[str, float] = {}
        self.avg_doc_length = 0.0
        self._indexed = False
    
    def add_documents(self, documents: List[Document]):
        """添加文档到BM25索引"""
        self.documents.extend(documents)
        
        # 构建词频统计
        for doc in documents:
            terms = self._tokenize(doc.content)
            term_freq = Counter(terms)
            self.doc_freqs.append(dict(term_freq))
        
        # 计算IDF值
        self._calculate_idf()
        
        # 计算平均文档长度
        total_length = sum(len(self._tokenize(doc.content)) for doc in self.documents)
        self.avg_doc_length = total_length / len(self.documents) if self.documents else 0
        
        self._indexed = True
    
    def _tokenize(self, text: str) -> List[str]:
        """分词"""
        # 简单的分词实现，实际应用中可以使用更复杂的分词器
        text = text.lower()
        # 保留中文、英文和数字
        text = re.sub(r'[^\u4e00-\u9fff\w\s]', ' ', text)
        tokens = text.split()
        return tokens
    
    def _calculate_idf(self):
        """计算IDF值"""
        doc_count = len(self.documents)
        term_doc_count = Counter()
        
        for doc_freq in self.doc_freqs:
            for term in doc_freq.keys():
                term_doc_count[term] += 1
        
        for term, doc_freq in term_doc_count.items():
            self.idf_values[term] = math.log((doc_count - doc_freq + 0.5) / (doc_freq + 0.5))
    
    def search(self, query: str, top_k: int = 10) -> List[SearchResult]:
        """BM25搜索"""
        if not self._indexed:
            return []
        
        query_terms = self._tokenize(query)
        scores = []
        
        for i, doc in enumerate(self.documents):
            score = self._calculate_bm25_score(query_terms, self.doc_freqs[i], len(self._tokenize(doc.content)))
            scores.append((i, score))
        
        # 按分数排序
        scores.sort(key=lambda x: x[1], reverse=True)
        
        results = []
        for doc_idx, score in scores[:top_k]:
            if score > 0:  # 只返回有得分的文档
                result = SearchResult(
                    document=self.documents[doc_idx],
                    score=score,
                    relevance_type="keyword"
                )
                results.append(result)
        
        return results
    
    def _calculate_bm25_score(self, query_terms: List[str], doc_term_freq: Dict[str, int], doc_length: int) -> float:
        """计算BM25分数"""
        score = 0.0
        
        for term in query_terms:
            if term in doc_term_freq:
                tf = doc_term_freq[term]
                idf = self.idf_values.get(term, 0)
                
                # BM25公式
                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (1 - self.b + self.b * (doc_length / self.avg_doc_length))
                score += idf * (numerator / denominator)
        
        return score


class HybridRetriever(RetrievalInterface, EventEmitter):
    """混合检索器 - 集成多种检索策略"""
    
    def __init__(self):
        super().__init__()
        self.data_manager: Optional[DataManager] = None
        self.embedding_generator: Optional[EmbeddingGenerator] = None
        self.bm25_retriever = BM25Retriever()
        self._initialized = False
        
        # 检索权重配置
        self.retrieval_weights = {
            "vector": 0.6,
            "keyword": 0.3,
            "graph": 0.1
        }
    
    async def initialize(self):
        """初始化混合检索器"""
        if self._initialized:
            return
        
        try:
            # 获取依赖组件
            self.data_manager = await container.resolve(DataManager)
            self.embedding_generator = await container.resolve(EmbeddingGenerator)
            
            await self.data_manager.initialize()
            await self.embedding_generator.initialize()
            
            # 初始化BM25索引（从向量数据库加载文档）
            await self._initialize_bm25_index()
            
            self._initialized = True
            
            await self.emit("hybrid_retriever_initialized", {
                "retrieval_weights": self.retrieval_weights
            })
            
        except Exception as e:
            await self.emit_error("hybrid_retriever_initialization", e)
            raise
    
    async def _initialize_bm25_index(self):
        """初始化BM25索引"""
        try:
            # 这里简化处理，实际应用中需要从数据库加载所有文档
            # 由于我们没有直接的方法获取所有文档，这里先跳过
            # 在实际应用中，应该在数据灌入时同步更新BM25索引
            pass
        except Exception as e:
            await self.emit_error("bm25_index_initialization", e)
    
    async def retrieve(
        self, 
        query: str, 
        top_k: int = 10,
        retrieval_modes: List[str] = None,
        metadata_filter: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """混合检索主方法"""
        if not self._initialized:
            await self.initialize()
        
        try:
            if retrieval_modes is None:
                retrieval_modes = ["vector", "keyword", "graph"]
            
            await self.emit("hybrid_retrieval_started", {
                "query": query[:100] + "..." if len(query) > 100 else query,
                "modes": retrieval_modes,
                "top_k": top_k
            })
            
            # 并行执行不同的检索策略
            retrieval_tasks = []
            
            if "vector" in retrieval_modes:
                retrieval_tasks.append(self._vector_retrieve(query, top_k, metadata_filter))
            
            if "keyword" in retrieval_modes:
                retrieval_tasks.append(self._keyword_retrieve(query, top_k))
            
            if "graph" in retrieval_modes:
                retrieval_tasks.append(self._graph_retrieve(query, top_k))
            
            # 等待所有检索任务完成
            retrieval_results = await asyncio.gather(*retrieval_tasks, return_exceptions=True)
            
            # 合并和融合结果
            all_results = []
            for i, result in enumerate(retrieval_results):
                if isinstance(result, Exception):
                    await self.emit_error(f"retrieval_mode_{retrieval_modes[i]}", result)
                    continue
                all_results.extend(result)
            
            # 结果融合
            fused_results = await self._fuse_results(all_results, retrieval_modes)
            
            # 排序和截断
            final_results = sorted(fused_results, key=lambda x: x.score, reverse=True)[:top_k]
            
            await self.emit("hybrid_retrieval_completed", {
                "query": query[:100] + "..." if len(query) > 100 else query,
                "total_results": len(all_results),
                "final_results": len(final_results),
                "modes_used": retrieval_modes
            })
            
            return final_results
            
        except Exception as e:
            await self.emit_error("hybrid_retrieve", e)
            return []
    
    async def _vector_retrieve(self, query: str, top_k: int, metadata_filter: Optional[Dict[str, Any]]) -> List[SearchResult]:
        """向量检索"""
        try:
            # 生成查询嵌入
            query_embedding = await self.embedding_generator.generate_query_embedding(query)
            
            # 执行向量搜索
            results = await self.data_manager.vector_store.search_similar(
                query_embedding, top_k * 2, metadata_filter  # 获取更多结果用于融合
            )
            
            return results
            
        except Exception as e:
            await self.emit_error("vector_retrieve", e)
            return []
    
    async def _keyword_retrieve(self, query: str, top_k: int) -> List[SearchResult]:
        """关键词检索"""
        try:
            # 使用BM25检索
            results = self.bm25_retriever.search(query, top_k * 2)
            return results
            
        except Exception as e:
            await self.emit_error("keyword_retrieve", e)
            return []
    
    async def _graph_retrieve(self, query: str, top_k: int) -> List[SearchResult]:
        """图检索"""
        try:
            # 基于实体的图检索
            results = await self.data_manager.hybrid_search(
                query, 
                search_modes=["graph"],
                top_k=top_k * 2
            )
            
            return results
            
        except Exception as e:
            await self.emit_error("graph_retrieve", e)
            return []
    
    async def _fuse_results(self, results: List[SearchResult], modes: List[str]) -> List[SearchResult]:
        """结果融合 - 使用RRF (Reciprocal Rank Fusion)"""
        try:
            # 按检索类型分组
            results_by_type = {}
            for result in results:
                result_type = result.relevance_type
                if result_type not in results_by_type:
                    results_by_type[result_type] = []
                results_by_type[result_type].append(result)
            
            # 为每种类型的结果排序
            for result_type in results_by_type:
                results_by_type[result_type].sort(key=lambda x: x.score, reverse=True)
            
            # 使用RRF融合
            doc_scores = {}
            k = 60  # RRF参数
            
            for result_type, type_results in results_by_type.items():
                weight = self.retrieval_weights.get(result_type, 0.1)
                
                for rank, result in enumerate(type_results):
                    doc_id = result.document.id
                    rrf_score = weight / (k + rank + 1)
                    
                    if doc_id not in doc_scores:
                        doc_scores[doc_id] = {
                            "score": 0.0,
                            "document": result.document,
                            "sources": []
                        }
                    
                    doc_scores[doc_id]["score"] += rrf_score
                    doc_scores[doc_id]["sources"].append({
                        "type": result_type,
                        "score": result.score,
                        "rank": rank + 1
                    })
            
            # 创建融合后的结果
            fused_results = []
            for doc_id, score_info in doc_scores.items():
                fused_result = SearchResult(
                    document=score_info["document"],
                    score=score_info["score"],
                    relevance_type="hybrid"
                )
                
                # 添加融合信息到元数据
                fused_result.document.metadata["fusion_info"] = {
                    "sources": score_info["sources"],
                    "fusion_method": "RRF"
                }
                
                fused_results.append(fused_result)
            
            return fused_results
            
        except Exception as e:
            await self.emit_error("fuse_results", e)
            return results  # 降级返回原始结果
    
    async def add_documents_to_index(self, documents: List[Document]):
        """向索引添加文档"""
        try:
            # 更新BM25索引
            self.bm25_retriever.add_documents(documents)
            
            await self.emit("documents_added_to_index", {
                "document_count": len(documents)
            })
            
        except Exception as e:
            await self.emit_error("add_documents_to_index", e)
    
    def update_retrieval_weights(self, weights: Dict[str, float]):
        """更新检索权重"""
        # 归一化权重
        total_weight = sum(weights.values())
        if total_weight > 0:
            self.retrieval_weights = {k: v / total_weight for k, v in weights.items()}
        
    async def explain_retrieval(self, query: str, doc_id: str) -> Dict[str, Any]:
        """解释检索结果"""
        try:
            explanation = {
                "query": query,
                "document_id": doc_id,
                "retrieval_breakdown": {}
            }
            
            # 分别计算各种检索方式的得分
            modes = ["vector", "keyword", "graph"]
            
            for mode in modes:
                try:
                    if mode == "vector":
                        results = await self._vector_retrieve(query, 100, None)
                    elif mode == "keyword":
                        results = await self._keyword_retrieve(query, 100)
                    elif mode == "graph":
                        results = await self._graph_retrieve(query, 100)
                    
                    # 查找目标文档
                    target_result = None
                    for result in results:
                        if result.document.id == doc_id:
                            target_result = result
                            break
                    
                    if target_result:
                        explanation["retrieval_breakdown"][mode] = {
                            "score": target_result.score,
                            "rank": results.index(target_result) + 1,
                            "weight": self.retrieval_weights.get(mode, 0.0)
                        }
                    else:
                        explanation["retrieval_breakdown"][mode] = {
                            "score": 0.0,
                            "rank": None,
                            "weight": self.retrieval_weights.get(mode, 0.0)
                        }
                        
                except Exception as e:
                    explanation["retrieval_breakdown"][mode] = {
                        "error": str(e)
                    }
            
            return explanation
            
        except Exception as e:
            await self.emit_error("explain_retrieval", e)
            return {"error": str(e)}